#include "qelib.h"
#include <pthread.h>
#include <winsock2.h>
#include <stdio.h>
#include <stdlib.h>

#pragma comment(lib, "ws2_32.lib")

QELOG_DOMAIN("telnet");

#define TELNET_PORT                     (23)
#define TELNET_MAX_CLIENTS              (4)

#define TELNET_CMD_WILL                 (0xFB)
#define TELNET_CMD_WONT                 (0xFC)
#define TELNET_CMD_DO                   (0xFD)

#define TELNET_OPTION_DO_ECHO           (01)
#define TELNET_OPTION_GO_AHEAD          (03)
#define TELNET_OPTION_TERMINAL_TYPE     (24)
#define TELNET_OPTION_NEG_WINDOW_SIZE   (31)

typedef struct
{
    pthread_t task;
    SOCKET socket;
    struct sockaddr_in addr;
    qe_shell shell;
    qe_gbuf *rxbuf;
    qe_list list;
} telnet_client;

typedef struct
{
    SOCKET socket;
    struct sockaddr_in addr;
    qe_uint num_clients;
    qe_list clients;
} telnet_server;

static telnet_client *tclient = QE_NULL;

static const char shell_show_info[] = 
    ""__DATE__" "__TIME__"\r\n";

static void telnet_log_output(const char *message)
{
    qe_printf("%s", message);
    qe_shell_write_end_line(&tclient->shell, (char *)message, qe_strlen(message));
}

static void telnet_send_response(telnet_client *tc, qe_u8 *response, int len)
{
    send(tc->socket, response, len, 0);
}

static void telnet_negotiation(telnet_client *tc, qe_u8 cmd, qe_u8 option)
{
    qe_u8 response[3];
    response[0] = 0xFF;

    switch (cmd) {

    case TELNET_OPTION_DO_ECHO:
        if (cmd == 0xFD) {
            response[1] = TELNET_CMD_WILL;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        }
        break;

    case TELNET_OPTION_GO_AHEAD:
        if (cmd == 0xFD) {
            response[1] = 0xFC;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        } else if (cmd == 0xFB) {
            response[1] = TELNET_CMD_DO;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        }
        break;

    case TELNET_OPTION_TERMINAL_TYPE:
        if (cmd == 0xFB) {
            response[1] = TELNET_CMD_DO;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        }
        break;

    case TELNET_OPTION_NEG_WINDOW_SIZE:
        if (cmd == 0xFB) {
            response[1] = TELNET_CMD_DO;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        }
        break;

    default:
        if (cmd == 0xFB || cmd == 0xFD) {
            response[1] = (cmd == 0xFB) ? 0xFE : 0xFC;
            response[2] = option;
            telnet_send_response(tc, response, sizeof(response));
        }
        break;
    }
}

static void telnet_input_process(telnet_client *tc)
{
    int i = 0;
    int len;
    qe_u8 *data;
    qe_u8 cmd;
    qe_u8 option;
    
    len = qe_gbuf_data_size(tc->rxbuf);
    data = (qe_u8 *)qe_gbuf_pos(tc->rxbuf);

    while (i < len) {

        /* IAC */
        if (data[i] == 0xFF) {
            if (i + 2 >= len)
                break;
            cmd = data[i+1];
            option = data[i+2];
            telnet_negotiation(tc, cmd, option);
            i+=3;
        } else {
            qe_shell_handler(&tc->shell, data[i]);
            i+=1;
        }
    }
}

static int ts_shell_read(qe_shell *sh, char *data, int len)
{
    return 0;
}

static int ts_shell_write(qe_shell *sh, char *data, int len)
{
    telnet_client *tc = tclient;
    if (!tc) {
        qe_error("tc null");
        return -1;
    }

    return send(tc->socket, data, len, 0);
}

static void *telnet_client_task(void *args)
{
    int n;
    telnet_client *tc = (telnet_client *)args;

    tc->rxbuf = qe_gbuf_new(1500);

    qe_shell_init(&tc->shell, 1024, CONFIG_SHELL_DEFAULT_USER, ts_shell_write, ts_shell_write);
    qe_debug("shell init");

    qe_shell_set_usr_info(&tc->shell, shell_show_info);
    qelog_set_default_output(telnet_log_output);

    qe_debug("shell init");

    while (1) {

        n = recv(tc->socket, qe_gbuf_pos(tc->rxbuf), qe_gbuf_size(tc->rxbuf), 0);
        qe_debug("recv:%d", n);
        if (n < 0) {
            qe_error("tc %p recv error:%d", tc->socket, n);
            closesocket(tc->socket);
            qe_list_remove(&tc->list);
            qe_free(tc);
            return QE_NULL;
        }
        qe_gbuf_data_size(tc->rxbuf) = n;

        qe_debug("recv n:%d", n);
        qe_hexdump_debug(qe_gbuf_pos(tc->rxbuf), n);

        telnet_input_process(tc);
    }
}

static void ts_new_connect(telnet_server *ts, SOCKET socket)
{
    telnet_client *tc;

    if (ts->num_clients >= TELNET_MAX_CLIENTS) {
        qe_warning("num client %d up to max, refuse connect", ts->num_clients);
        closesocket(socket);
        return;
    }

    tc = qe_malloc(sizeof(telnet_client));
    qe_assert(tc);
    qe_memset(tc, 0, sizeof(telnet_client));

    tc->socket = socket;
    tclient = tc;

    qe_info("add new client");

    pthread_create(&tc->task, QE_NULL, telnet_client_task, tc);
}

static void *loop_task(void *args)
{
    while (1) {
        qe_debug("loop task");
        qe_sleep(1);
    }
}

int main(int argc, char *argv[])
{
    int addrlen;
    SOCKET client;
    pthread_t tid;
    telnet_server *ts;
    struct sockaddr_in client_addr;

    qelog_init(QELOG_DEBUG, QELOG_CL|QELOG_DM|QELOG_DATE);

    ts = qe_malloc(sizeof(telnet_server));
    qe_assert(ts);

    qe_memset(ts, 0, sizeof(telnet_server));

    qe_list_init(&ts->clients);

    qe_info("telnet server start");

    /* Initialize winsock */
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        qe_error("WSAStartup failed: %d", WSAGetLastError());
        return -1;
    }

    /* Create server socket */
    ts->socket = socket(AF_INET, SOCK_STREAM, 0);
    qe_assert(ts->socket != INVALID_SOCKET);

    /* Setup server address */
    addrlen = sizeof(client_addr);
    ts->addr.sin_family = AF_INET;
    ts->addr.sin_port = qe_htons(TELNET_PORT);
    ts->addr.sin_addr.s_addr = INADDR_ANY;

    /* Bind server address */
    if (bind(ts->socket, (struct sockaddr *)&ts->addr, sizeof(ts->addr)) == SOCKET_ERROR) {
        qe_error("bind server addr failed");
        closesocket(ts->socket);
        WSACleanup();
        return -1;
    }

    /* Start listen */
    if (listen(ts->socket, TELNET_MAX_CLIENTS) == SOCKET_ERROR) {
        qe_error("listen failed");
        closesocket(ts->socket);
        WSACleanup();
        return -1;
    }

    qe_info("start listening");

    pthread_create(&tid, QE_NULL, loop_task, QE_NULL);

    while (1) {
        client = accept(ts->socket, (struct sockaddr *)&client_addr, &addrlen);
        if (client != INVALID_SOCKET) {
            qe_info("accept %p", client);
            ts_new_connect(ts, client);
        }
        qe_msleep(1000);
    }

    return 0;
}
